#!/usr/bin/env python
    
import sys
import pwd
import re
import os
import getpass
import os.path

from twisted.internet.protocol import Factory
from twisted.internet.endpoints import TCP4ClientEndpoint
from twisted.internet import reactor

this_dir  = os.path.dirname( os.path.abspath(__file__) )
build_dir = os.path.join(this_dir, 'build')

if os.path.exists(build_dir):
    for p in (os.path.join(build_dir, fn) for fn in os.listdir(build_dir) if fn.startswith('lib.')):
        sys.path.insert(0,p)

import srp

from srp_auth import client


result = 1

sha_map = dict( sha1   = srp.SHA1,
                sha224 = srp.SHA224,
                sha256 = srp.SHA256,
                sha384 = srp.SHA384,
                sha512 = srp.SHA512 )
    
prime_map = { 1024 : srp.NG_1024, 2048 : srp.NG_2048,
              4096 : srp.NG_4096, 8192 : srp.NG_8192 }
    

cache_table_def = '''create table cache (username text UNIQUE PRIMARY KEY,
                                         salt blob NOT NULL,
                                         verifier blob NOT NULL,
                                         timestamp integer NOT NULL)'''

root_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

default_config_file = os.path.join(root_dir, 'etc', 'srp_auth.conf')
default_key_file    = os.path.join(root_dir, 'var', 'lib', 'srp_auth', 'default_key.der')
default_cred_cache  = os.path.join(root_dir, 'var', 'lib', 'srp_auth', 'cred_cache.db')
default_cred_db     = os.path.join(root_dir, 'var', 'lib', 'srp_auth', 'credentials.db')


usage_string = '''
Usage: srp_admin [generic options] add username [-l] [-a]
       srp_admin [generic options] remove username
       srp_admin [generic options] list
       srp_admin [generic options] show username
       srp_admin [generic options] change_password [username]
       srp_admin [generic options] lock username
       srp_admin [generic options] unlock username
       srp_admin [generic options] set_admin username
       srp_admin [generic options] unset_admin username
       srp_admin [generic options] set_expiration username time-specification
       srp_admin get_public_info [server:port] [file_name]
       srp_admin init_host <server>:<port> [-c <num_cached_credentials>]
       srp_admin create_credential_cache [file_name]
       srp_admin create_credential_database [file_name] [-h <hash_type>] [-p <prime_size> ]

srp_admin - Administration tool for SRP Authentication databases

Generic options:
  -u <username> Specifies the user to connect as. The password for this user will
                user will be queried interactively.
  -f <filename> Configuration file. Default = %(root)s/etc/srp_auth.conf
  -c <cfg_name> Use the specified configuration instead of the default

Commands:

  add                        - Adds a user to the database. If '-l' is present,
                               the account will be initially locked. If '-a'
                               is present, the account will have administrative
                               privileges.
  remove                     - Removes a user from the database
  list                       - List all user names
  show                       - Display user properties
  change_password            - Changes the users password (defaults to
                               current user if username is not specified)
  lock                       - Locks the specified user account
  unlock                     - Unlocks the specified user account
  set_admin                  - Allows the specified user to administer the
                               database
  unset_admin                - Prevents the specified user from administering
                               the database
  set_expiration             - Sets the expiration time of the account
                               See formatting notes below.
  get_public_info            - Retrieves the server's required hash algorithm and
                               prime number sizes. If a filename is provided, this
                               command will store the server's binary, DER-encoded
                               public RSA key in the file. If no server:port is
                               specified, the server will be read from the
                               configuration file.
  init_host                  - Creates the %(root)s/etc/srp_auth.conf and 
                               associated files for using the specified SRP server. 
  create_credential_cache    - Creates a creditial cache for authentication
                               while disconnected from the SRP server
  create_credential_database - Creates a credential database for an SRP
                               server. The optional hash and prime arguments may
                               be used to tune the security/performance tradeoff
                               of the Secure Remote Password protocol. The defaults
                               are SHA1 and 2048-bits. Additional options are
                               listed in the SRP Tuning section.

Account Expiration Time Formatting:
  now             - Sets account to expired state
  never           - Account never expires

  <number><units> - Sets the account to expire a number of units from
                    now. Acceptable units are:
                      s - Seconds
                      n - Minutes (note this is lower case N)
                      h - Hours
                      d - Days
                      m - Months (note this is lower case M)
                      y - Years

  Examples:
    3d   - 3 days from now
    5n   - 5 Minutes from now


Host Initialization:

  The 'init_host' command provides a simple means to configure the local host
  for use with the specified SRP server. This command is purely optional and
  serves only to automate the otherwise manual process of setting up a default
  srp_auth configuration.

  The 'init_host' command will create an %(root)s/etc/srp_auth.conf file
  with the proper settings for use with the specified SRP server and places
  the SRP server's public key in the %(root)s/etc/srp_auth.d/default.der 
  file. It will also create a credential cache file in the
  %(root)s/var/lib/srp_auth directory. The -c option may be used to
  override the default number of cached credentials used by the PAM plugin
  module (defaults to 10). Specifying a value of 0 will disable the use of
  the credential cache.
  

  This command does *NOT* modify the PAM configuration files. The PAM
  configuration must be manually edited for proper use of the srp_auth module.
  Please refer to the srp_auth(TODO) manual page for instructions on
  modifying the PAM configuration to use the srp_auth plugin module.
  

SRP Tuning:
  The credential database used for verifying user passwords can be configured
  to use multiple length hash algorithms and prime numbers. Essentially, longer
  hashes and prime numbers result in greater security at the cost of increased
  computation time. The defaults of SHA1 hashes and 2048-bit prime numbers strike
  a good balance. However, these parameters may be tuned as follows:

  Available hash types:
      sha1, sha224, sha256, sha384, sha512

  Available prime number sizes:
      1024, 2048, 4096, 8192
'''


config_tmpl = '''
server          = %(server)s
port            = %(port)d
public_key_file = %(key_file)s
hash_algorithm  = %(hash_alg)s
prime_size      = %(prime_size)d

%(cache)s
'''


def err_exit( msg ):
    print >> sys.stderr, msg
    sys.exit(1)


def usage():
    err_exit( usage_string % { 'root' : root_dir } )

        
def get_password( username, confirm = False ):
    pw1 = ''
    pw2 = 'x'
    while pw1 != pw2:
        pw1 = getpass.getpass('Enter password for %s:' % username)
        if not confirm:
            return pw1
        pw2 = getpass.getpass('Confirm password for %s:' % username)
        if not pw1 == pw2:
            print 'Passwords do not match\n'
    return pw1


def parse_config( cfg_file ):
    cfg = dict()
    sect = dict()
    cfg['default'] = sect
    sect_re = re.compile('\s*\[.+]\s*$')
    
    with open(cfg_file, 'r') as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith('#'):
                continue
            m = sect_re.match(line)
            if m:
                sect_name = line.strip()[1:-1].strip()
                sect = dict()
                cfg[ sect_name ] = sect

            elif '=' in line:
                k,v = line.split('=')
                sect[ k.strip() ] = v.strip()

    return cfg


def secure_create_file( filename, mode, data = None ):
    fd = None
    try:
        fd = os.open(filename, os.O_WRONLY | os.O_EXCL | os.O_CREAT, mode)
    except OSError, err:
        raise Exception('Secure creation of file failed: %s' % str(err))
    if data:
        os.write(fd, data)
    os.close(fd)
    
        


class Options (object):
    def __init__(self):
        
        self.args        = sys.argv[1:]

        # --- Generic Arguments ---
        self.username    = self.extract_param( '-u', pwd.getpwuid(os.getuid())[0] )
        self.config_file = self.extract_param( '-f', default_config_file )
        self.srp_config  = self.extract_param( '-c', 'default' )
        
        try:
            self.command     = self.args.pop(0)
        except IndexError:
            usage()

        self.server     = None
        self.port       = None
        self.hash_type  = srp.SHA1
        self.prime_type = srp.NG_2048

        try:
            cmd_method = getattr(self, 'parse_' + self.command)
        except AttributeError:
            usage()

        cmd_method()

        # Each argument parsing method must consume all provided arguments. If there's anything
        # left, the user provided extra, erroneous arguments.
        if self.args:
            print 'Invalid arguments'
            usage()

        
    def extract_param(self, flag, default = None):
        try:
            i = self.args.index(flag)
            t = self.args[ i + 1 ]
            del self.args[i:i+2]
            return t
        except:
            if default:
                return default
            else:
                usage()


    def extract_bool(self, flag, on_present = True):
        try:
            self.args.remove( flag )
            return on_present
        except:
            return not on_present


    def next_arg(self):
        try:
            return self.args.pop(0)
        except IndexError:
            usage()


    def assert_no_file(self, fn = None ):
        if fn is None:
            fn = self.file_name
        if os.path.exists( fn ):
            err_exit( 'Failure. Target file already exists: ' + fn )


    def read_config_file(self):
        if not os.path.exists( self.config_file ):
            err_exit( 'Missing configuration file: ' + self.config_file)
    
        try:
            cfile = parse_config( self.config_file )
        except:
            err_exit( 'Failed to parse config file: ' + self.config_file )

        try:
            cfg = cfile[ self.srp_config ]
        except KeyError:
            err_exit( 'Invalid configuration file section: ' + self.srp_config )

        try:
            self.server = cfg['server']
        except KeyError:
            err_exit( 'Missing "server" definition in %s section of the configuration file' % self.srp_config )

        try:
            self.port = int( cfg['port'] )
        except KeyError:
            err_exit( 'Missing "port" definition in %s section of the configuration file' % self.srp_config )
        except ValueError:
            err_exit( 'Invalid "port" definition in %s section of the configuration file' % self.srp_config )
    
        if 'hash_algorithm' in cfg:
            try:
                self.hash_type = sha_map[ cfg['hash_algorithm'].lower() ]
            except KeyError:
                err_exit( 'Invalid "hash_algorithm" definition in %s section of the configuration file' % self.srp_config )

        if 'prime_size' in cfg:
            try:
                self.prime_type = prime_map[ int(cfg['prime_size']) ]
            except KeyError:
                err_exit( 'Invalid "prime_size" definition in %s section of the configuration file' % self.srp_config )
            except ValueError:
                err_exit( 'Invalid "prime_size" definition in %s section of the configuration file' % self.srp_config )


            
    # ---- Command Specific Argument Processors ----

    def parse_add(self):
        self.is_locked       = self.extract_bool('-l')
        self.is_admin        = self.extract_bool('-a')
        self.expires_at      = -1 # never
        self.target_username = self.next_arg()

        if not re.match('^[_a-zA-Z][_a-zA-Z0-9.]+$', self.target_username):
            err_exit("Invalid username: " + self.target_username)

        self.new_password = get_password( self.target_username, True )
        


    def parse_remove(self):
        self.target_username = self.next_arg()

    
    def parse_list(self):
        pass

    
    def parse_show(self):
        self.target_username = self.next_arg()

        
    def parse_change_password(self):
        self.target_username = self.username if not self.args else self.next_arg()
        self.new_password    = get_password( self.target_username, True )

        
    def parse_lock(self):
        self.target_username = self.next_arg()

        
    def parse_unlock(self):
        self.target_username = self.next_arg()

        
    def parse_set_admin(self):
        self.target_username = self.next_arg()

        
    def parse_unset_admin(self):
        self.target_username = self.next_arg()

        
    def parse_set_expiration(self):
        self.target_username = self.next_arg()

        when = self.next_arg().lower()

        try:
            if when == 'now':
                self.expires_at = 0
                
            elif when == 'never':
                self.expires_at = -1
                
            else:
                num        = when[:-1]
                time_key   = when[-1:].lower()
            
                factor_map = dict( s = 1, n = 60, h = 60*60, d=60*60*24,
                                   m=60*60*24*30, y=60*60*24*30*12 )
                
                num        = int(num)
                
                self.expires_at = num * factor_map[ time_key ]
                
        except:
            usage()

            
    def parse_get_public_info(self):
        self.file_name = None
        
        if len(self.args) > 2:
            usage()
            
        while self.args:
            arg = self.next_arg()

            m = re.match('(.+):(\d+)', arg)

            if m:
                self.server   = m.group(1)
                try:
                    self.port = int( m.group(2) )
                except ValueError:
                    err_exit("Invalid port number: " + m.group(2))
            else:
                self.file_name = arg
                self.assert_no_file()


    def parse_init_host(self):
        self.ncache = self.extract_param( '-c', '10' )
        try:
            self.ncache = int(self.ncache)
        except ValueError:
            err_exit('Invalid number: %s' % self.ncache)

        self.assert_no_file( default_config_file )
        self.assert_no_file( default_key_file    )

        if self.ncache != 0:
            self.assert_no_file( default_cred_cache )

        svr_string = self.next_arg()

        m = re.match('(.+):(\d+)', svr_string)

        if not m:
            err_exit('Invalid username@server:port specification: ' + svr_string)

        self.server   = m.group(1)
        try:
            self.port = int(m.group(2))
        except ValueError:
            err_exit("Invalid port number: " + m.group(2))
        

    def parse_create_credential_cache(self):
        if self.args:
            self.file_name = self.next_arg()
        else:
            self.file_name = default_cred_cache
        self.assert_no_file()

        
    def parse_create_credential_database(self):
        try:
            self.hash_type = sha_map[ self.extract_param( '-h', 'sha1' ).lower() ]
        except:
            err_exit('Invalid hash type')

        try:
            self.prime_type = prime_map[ int(self.extract_param( '-p', '2048' )) ]
        except:
            err_exit('Invalid prime number')

        if self.args:
            self.file_name = self.next_arg()
        else:
            self.file_name = default_cred_db
        self.assert_no_file()
        


#-------------------------------------------------
# Non-Neworked Commands
#
def create_credential_cache( ops ):
    import sqlite3

    if os.path.exists( ops.file_name ):
        print >> sys.stderr, 'File already exists: %s' % ops.file_name
        sys.exit(1)

    secure_create_file( ops.file_name, 0660 )

    con  = sqlite3.connect(ops.file_name)
    c    = con.cursor()
    c.execute( cache_table_def )
    con.commit()
    con.close()





def create_credential_database( ops ):
    import getpass
    from   srp_auth import db

    admin_name     = getpass.getuser()
    adname         = admin_name
    admin_password = ''
    
    adname = raw_input('Enter primary administrator username [%s]: ' % adname)
    adname = adname.strip()
        
    if not adname:
        adname = admin_name
        
    admin_password = get_password( adname, True )
        
    s = db.DB( ops.file_name, True, ops.hash_type, ops.prime_type )

    salt, ver = srp.create_salted_verification_key( admin_name, admin_password, hash_alg=ops.hash_type, ng_type=ops.prime_type )
    
    s.add_user( admin_name, salt, ver, is_admin=True )
    s.close()

#-------------------------------------------------
# Neworked Commands
#
def connect_failed(err):
    global result
    result = 1
    print >> sys.stderr, 'Connect Failed: ', err.getErrorMessage()
    reactor.stop()

def cmd_fail(err):
    global result
    result = 1
    print >> sys.stderr, 'Command Failed: ', err.getErrorMessage()
    reactor.stop()

    
def on_connected(ops, srp_client):

    def command_ok( arg ):
        global result
        result = 0
        reactor.stop()
    
    
    def command_fail( err ):
        global result
        result = 1
        print >> sys.stderr, err.getErrorMessage()
        reactor.stop()

        
    if ops.command == 'add':
        d = srp_client.add_user(ops.target_username, ops.new_password, ops.is_locked, ops.expires_at, ops.is_admin)
        
    elif ops.command == 'remove':
        d = srp_client.remove_user(ops.target_username)
        
    elif ops.command == 'list':
        d = srp_client.get_users()
        
        def gotem( ulist ):
            for u in ulist:
                print u
                
        d.addCallback( gotem )
        
    elif ops.command == 'show':
        d = srp_client.get_user_properties( ops.target_username )
        
        def gotit( p ):
            print 'Username: ', p.username
            print 'Locked  : ', 'True' if p.is_locked else 'False'
            print 'Admin   : ', 'True' if p.is_admin  else 'False'
            print 'Expires : ', p.expires_at
            
        d.addCallback( gotit )
        
    elif ops.command == 'change_password':
        d = srp_client.change_password( ops.target_username, ops.new_password )
        
    elif ops.command == 'lock':
        d = srp_client.set_user_lock( ops.target_username, True )
        
    elif ops.command == 'unlock':
        d = srp_client.set_user_lock( ops.target_username, False )
        
    elif ops.command == 'set_admin':
        d = srp_client.set_user_admin( ops.target_username, True )
        
    elif ops.command == 'unset_admin':
        d = srp_client.set_user_admin( ops.target_username, False )
        
    elif ops.command == 'set_expiration':
        d = srp_client.set_user_expiration( ops.target_username, ops.expires_at )
        
    d.addCallbacks( command_ok, command_fail )
    
    

def get_public_info( ops ):
    def got_it( tpl ):
        
        hash_alg, prime_size, pub_key = tpl

        hmap = { 0 : 'SHA1', 1 : 'SHA224', 2 : 'SHA256', 3 : 'SHA384', 4 : '512' }
        pmap = { 0 : 1024, 1 : 2048, 2 : 4096, 3 : 8192 }
        
        print 'Hash Algorithm:', hmap[ hash_alg ]
        print 'Prime Size:', pmap[prime_size]

        if ops.file_name is not None:
            try:
                secure_create_file( ops.file_name, 0644, data = pub_key )
            except:
                raise Exception('Failed to create public key file. File already exists')

        reactor.stop()
    
    if not ops.server:
        ops.read_config_file()

    d = client.fetchSRPServerPublicInfo( ops.server, ops.port )
    d.addCallbacks( got_it )
    d.addErrback( cmd_fail )

    reactor.run()


def init_host( ops ):
    
    def got_it( tpl ):
        hash_alg, prime_size, pub_key = tpl

        key_dir   = os.path.dirname( default_key_file )
        cache_dir = os.path.dirname( default_cred_cache )
        
        hmap = { 0 : 'SHA1', 1 : 'SHA224', 2 : 'SHA256', 3 : 'SHA384', 4 : 'SHA512' }
        pmap = { 0 : 1024, 1 : 2048, 2 : 4096, 3 : 8192 }

        cache_str = ''
        if ops.ncache != 0:
            cache_str = 'credential_cache = %s\nmax_cache_entries = %d\n' % ( default_cred_cache, ops.ncache )
            if not os.path.isdir( cache_dir ):
                try:
                    os.makedirs( cache_dir, 0750 )
                except OSError, e:
                    raise Exception('Failed to create directory: %s. Reason: %s' % (key_dir, e.strerror))


        cfg_data = config_tmpl % { 'server'     : ops.server,
                                   'port'       : ops.port,
                                   'key_file'   : default_key_file,
                                   'hash_alg'   : hmap[hash_alg],
                                   'prime_size' : pmap[prime_size],
                                   'cache'      : cache_str } 

        if not os.path.isdir( key_dir ):
            try:
                os.makedirs( key_dir, 0755 )
            except OSError, e:
                raise Exception('Failed to create directory: %s. Reason: %s' % (key_dir, e.strerror))
        
        try:
            secure_create_file( default_key_file, 0644, data = pub_key )
        except OSError, e:
            raise Exception('Failed to create public key file: ' + e.strerror)

        try:
            secure_create_file( default_config_file, 0644, data = cfg_data )
        except OSError, e:
            raise Exception('Failed to create the configuration file: ' + e.strerror)

        if ops.ncache != 0:
            ops.file_name = default_cred_cache
            try:
                create_credential_cache( ops )
            except OSError, e:
                raise Exception('Failed to create credential cache file: ' + e.strerror)

        reactor.stop()
            

    d = client.fetchSRPServerPublicInfo( ops.server, ops.port )
    d.addCallbacks( got_it )
    d.addErrback( cmd_fail )

    reactor.run()
    
    

def network_command( ops ):

    # All network commands require the config file and a valid
    # username/password
    ops.read_config_file()
    
    password = get_password(ops.username)

        
    def connected_cb( srp_client ):
        return on_connected( ops, srp_client )

    f = client.SrpClientFactory( ops.username, password, ops.hash_type, ops.prime_type )

    f.d_connected.addCallbacks( connected_cb, connect_failed )

    ep = TCP4ClientEndpoint(reactor, ops.server, ops.port)

    ep.connect( f ).addErrback( connect_failed )
    
    reactor.run()

    sys.exit(result)


#---------------------------------------------------------------------------------------------------
# "Main"
#
ops = Options()

if ops.command == 'create_credential_cache':
    
    create_credential_cache( ops )
        
elif ops.command == 'create_credential_database':
    create_credential_database( ops )

elif ops.command == 'get_public_info':
    get_public_info( ops )

elif ops.command == 'init_host':
    init_host( ops )

else:
    network_command( ops )


